from abc import abstractmethod
from algorithms.inferencer_base import InferencerBase


class InferencerBaseGenerate(InferencerBase):
    def __init__(self, model_name, device='cuda'):
        super().__init__(model_name=model_name, device=device)
        self.device = device
        self.model = self._load_model(model_name)

    @abstractmethod
    def _load_model(self, model_name):
        """加载模型，需由子类实现"""
        pass

    @abstractmethod
    def preprocess(self, input_data):
        """对输入数据进行预处理"""
        pass

    @abstractmethod
    def postprocess(self, outputs):
        """对模型输出进行后处理"""
        pass

    def inference(self, input_data):
        """标准推理流程：预处理 -> 推理 -> 后处理"""
        preprocessed = self.preprocess(input_data)
        raw_output = self.model(preprocessed)
        result = self.postprocess(raw_output)
        return result

    def inference_img2img(self, input_img):
        """图像到图像的推理流程"""
        return self.inference(input_img)

    def inference_txt2img(self, input_txt):
        """文本生成图像的推理流程"""
        raise NotImplementedError("该方法尚未实现")

    def inference_txtimg2img(self, input_txt, input_img):
        """图文联合生成图像的推理流程"""
        raise NotImplementedError("该方法尚未实现")